# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openFuyao is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#          http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

"""
This module defines the NPUDeviceInfoCollector class,
which is responsible for collecting information about the NPU devices,
including details about the cards and devices in the system. It interacts with 
the NPU utility functions to gather the device information and stores the results in a cache.
"""

import logging
import os

from hardware_monitor.base.base_device_info_collector import BaseDeviceInfoCollector
from hardware_monitor.collector.metrics.npu.npu_utils import get_card_list, get_device_num_in_card, get_device_logic_id

module_name = os.getenv("MODULE_NAME", "default")

logger = logging.getLogger(module_name)


class NPUDeviceInfoCollector(BaseDeviceInfoCollector):
    """
    Collects information about NPU devices, including card IDs, device IDs, and logic IDs.
    """
    def collect_device_info(self):
        self.cache.clear_all()

        card_count, card_list, _, error = get_card_list()
        if error or card_count == 0:
            logger.error("[NPUDeviceInfoCollector] %s", error)
            return

        for card_id in card_list:
            device_count, error = get_device_num_in_card(card_id)
            if error:
                logger.error("[NPUDeviceInfoCollector] %s", error)
                continue

            for device_id in range(device_count):
                logic_id, error = get_device_logic_id(card_id, device_id)
                if error:
                    logger.error("[NPUDeviceInfoCollector] %s", error)
                    continue

                device_info = {
                    "card_id": card_id,
                    "device_id": device_id,
                    "logic_id": logic_id
                }

                self.update_cache(device_info)
